#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time    : 2025/6/17 21:37
# @Author  : wuyupeng
# @FileName: gpg_key_manager.py
# @Software: PyCharm

import os
import shutil
import tempfile
import gnupg


class GPGKeyManager:
    def __init__(self, gnupghome=None):
        """
        Initialize GPG environment.
        :param gnupghome: Optional path to GnuPG home directory
                          (creates temporary directory if None)
        """
        self.gnupghome = gnupghome
        self._temp_dir = None
        self.gpg = None
        
        if not self.gnupghome:
            self._temp_dir = tempfile.mkdtemp()
            self.gnupghome = self._temp_dir
        
        os.makedirs(self.gnupghome, exist_ok=True)
        self.gpg = gnupg.GPG(gnupghome=self.gnupghome)
        self.gpg.encoding = 'utf-8'
    
    def __del__(self):
        """Clean up temporary directory on destruction"""
        if self._temp_dir and os.path.exists(self._temp_dir):
            shutil.rmtree(self._temp_dir)

    def generate_key(self, name, email, passphrase=None, key_type='RSA', key_length=4096):
        """
        Generate a new GPG key pair.
        Returns key fingerprint.
        """
        input_data = self.gpg.gen_key_input(
            key_type=key_type,
            key_length=key_length,
            name_real=name,
            name_email=email,
            passphrase=passphrase
        )
        key = self.gpg.gen_key(input_data)
        return key.fingerprint

    def import_key(self, key_data, passphrase=None):
        """
        Import existing private key into keyring.
        Returns key fingerprint.
        """
        result = self.gpg.import_keys(key_data)
        if not result.results:
            raise ValueError("No keys found in import data")
        
        # Extract fingerprint from import results
        fingerprint = result.results[0].get('fingerprint')
        
        # If key has passphrase, we need to unlock it
        if passphrase:
            self.gpg.trust_keys(fingerprint, 'TRUST_ULTIMATE')
            unlock_result = self.gpg.import_keys(key_data, passphrase=passphrase)
            if not unlock_result:
                raise RuntimeError("Failed to unlock key with passphrase")
                
        return fingerprint
    
    def export_public_key(self, fingerprint):
        """Export public key string"""
        return self.gpg.export_keys(fingerprint)

    def export_private_key(self, fingerprint, passphrase):
        """Export private key string"""
        return self.gpg.export_keys(fingerprint, secret=True, passphrase=passphrase)

    def sign_detached(self, file_path, fingerprint, passphrase=None, output_path=None):
        """
        Create detached signature (Release.gpg format)
        :return: Path to signature file
        """
        output_path = output_path or file_path + '.gpg'
        
        with open(file_path, 'rb') as f:
            status = self.gpg.sign_file(
                f,
                keyid=self.get_key(fingerprint)['keyid'],
                passphrase=passphrase,
                detach=True,
                output=output_path
            )
        
        if not status:
            raise RuntimeError(f"Detached signing failed: {status.status}")
        return output_path

    def sign_clearsigned(self, file_path, fingerprint, passphrase=None, output_path=None):
        """
        Create clearsigned signature (InRelease format)
        :return: Path to signed file
        """
        output_path = output_path or file_path + '.inrelease'
        
        with open(file_path, 'rb') as f:
            signed_data = self.gpg.sign_file(
                f,
                keyid=self.get_key(fingerprint)['keyid'],
                passphrase=passphrase,
                clearsign=True
            )
        
        if not signed_data:
            raise RuntimeError("Clearsign signing failed")
        
        with open(output_path, 'wb') as f_out:
            f_out.write(signed_data.data)
            
        return output_path

    def verify_detached(self, file_path, signature_path):
        """
        Verify a detached signature (Release.gpg)
        :return: Verification object
        """
        with open(signature_path, 'rb') as sig_file:
            return self.gpg.verify_file(sig_file, file_path)

    def verify_clearsigned(self, signed_file_path):
        """
        Verify a clearsigned file (InRelease)
        :return: Verification object
        """
        with open(signed_file_path, 'rb') as f:
            return self.gpg.verify_file(f)

    def get_key(self, fingerprint):
        """Get key information by fingerprint"""
        keys = self.gpg.list_keys(keys=fingerprint)
        return keys[0] if keys else None

    @property
    def keys(self):
        """List all available keys in keyring"""
        return self.gpg.list_keys()
    
